//
//  BAAlbumListLayout.m
//  BabyAlbum
//
//  Created by master on 17/2/14.
//  Copyright © 2017年 master. All rights reserved.
//

#import "BAAlbumListLayout.h"

@implementation BAAlbumListLayout {
    NSMutableArray *_columnHeights;
    NSMutableArray *_itemAttributes;
    NSMutableDictionary *_headerAttributes;
    NSMutableDictionary *_footerAttributes;

    NSMutableArray *_allAttributes;
}

- (instancetype)init
{
    self = [super init];
    if (self != nil) {
        self.numberOfColumns = 3;
        self.columnSpacing = 10.0;
        self.sectionInset = UIEdgeInsetsMake(10.0, 10.0, 10.0, 10.0);
        self.interItemSpacing = UIEdgeInsetsMake(10.0, 0, 10.0, 0);
    }
    return self;
}

- (void)prepareLayout
{
    _itemAttributes = [NSMutableArray array];
    _columnHeights = [NSMutableArray array];
    _allAttributes = [NSMutableArray array];
    _headerAttributes = [NSMutableDictionary dictionary];
    _footerAttributes = [NSMutableDictionary dictionary];

    CGFloat top = 0;
    
    NSInteger numberOfSections = [self.collectionView numberOfSections];
    for (NSUInteger section = 0; section < numberOfSections; section++) {
        NSInteger numberOfItems = [self.collectionView numberOfItemsInSection:section];
        
        top += _sectionInset.top;
        
        if (_headerHeight > 0) {
            CGSize headerSize = [self headerSizeForSection:section];
            UICollectionViewLayoutAttributes *attributes = [UICollectionViewLayoutAttributes
                                                            layoutAttributesForSupplementaryViewOfKind:UICollectionElementKindSectionHeader
                                                            withIndexPath:[NSIndexPath indexPathForItem:0 inSection:section]];
            attributes.frame = CGRectMake(_sectionInset.left, top, headerSize.width, headerSize.height);
            _headerAttributes[@(section)] = attributes;
            [_allAttributes addObject:attributes];
            top = CGRectGetMaxY(attributes.frame);
        }
        

        
        
        [_columnHeights addObject:[NSMutableArray array]];
        for (NSUInteger idx = 0; idx < self.numberOfColumns; idx++) {
            [_columnHeights[section] addObject:@(top)];
        }
        
        CGFloat columnWidth = [self _columnWidthForSection:section];
        [_itemAttributes addObject:[NSMutableArray array]];
        for (NSUInteger idx = 0; idx < numberOfItems; idx++) {
            NSUInteger columnIndex = [self _shortestColumnIndexInSection:section];
            NSIndexPath *indexPath = [NSIndexPath indexPathForItem:idx inSection:section];
            
            CGSize itemSize = [self itemSizeAtIndexPath:indexPath];
            CGFloat xOffset = _sectionInset.left + (columnWidth + _columnSpacing) * columnIndex;
            CGFloat yOffset = [_columnHeights[section][columnIndex] floatValue];
            
            UICollectionViewLayoutAttributes *attributes = [UICollectionViewLayoutAttributes
                                                            layoutAttributesForCellWithIndexPath:indexPath];
            attributes.frame = CGRectMake(xOffset, yOffset, itemSize.width, itemSize.height);
            
            _columnHeights[section][columnIndex] = @(CGRectGetMaxY(attributes.frame) + _interItemSpacing.bottom);
            
            [_itemAttributes[section] addObject:attributes];
            [_allAttributes addObject:attributes];
        }
        
        NSUInteger columnIndex = [self _tallestColumnIndexInSection:section];
        top = [_columnHeights[section][columnIndex] floatValue] - _interItemSpacing.bottom + _sectionInset.bottom;
        
        for (NSUInteger idx = 0; idx < [_columnHeights[section] count]; idx++) {
            _columnHeights[section][idx] = @(top);
        }
        
        if (_footerHeight > 0) {
            CGSize footerSize = [self footerSizeForSection:section];
            UICollectionViewLayoutAttributes *attributes = [UICollectionViewLayoutAttributes
                                                            layoutAttributesForSupplementaryViewOfKind:UICollectionElementKindSectionFooter
                                                            withIndexPath:[NSIndexPath indexPathForItem:0 inSection:section]];
            attributes.frame = CGRectMake(_sectionInset.left, top, footerSize.width, footerSize.height);
            _footerAttributes[@(section)] = attributes;
            [_allAttributes addObject:attributes];
            top = CGRectGetMaxY(attributes.frame);
        }
    }
}

- (NSArray *)layoutAttributesForElementsInRect:(CGRect)rect
{
    NSMutableArray *includedAttributes = [NSMutableArray array];
    // Slow search for small batches
    for (UICollectionViewLayoutAttributes *attributes in _allAttributes) {
        if (CGRectIntersectsRect(attributes.frame, rect)) {
            [includedAttributes addObject:attributes];
        }
    }
    return includedAttributes;
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
{
    if (indexPath.section >= _itemAttributes.count) {
        return nil;
    } else if (indexPath.item >= [_itemAttributes[indexPath.section] count]) {
        return nil;
    }
    return _itemAttributes[indexPath.section][indexPath.item];
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryViewOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
{
    if ([elementKind isEqualToString:UICollectionElementKindSectionHeader]) {
        return _headerAttributes[@(indexPath.section)];
    } else if([elementKind isEqualToString:UICollectionElementKindSectionFooter]) {
        return _footerAttributes[@(indexPath.section)];
    }
    return nil;
}

- (BOOL)shouldInvalidateLayoutForBoundsChange:(CGRect)newBounds
{
    if (!CGSizeEqualToSize(self.collectionView.bounds.size, newBounds.size)) {
        return YES;
    }
    return NO;
}

- (CGFloat)_widthForSection:(NSUInteger)section
{
    return self.collectionView.bounds.size.width - _sectionInset.left - _sectionInset.right;
}

- (CGFloat)_columnWidthForSection:(NSUInteger)section
{
    return ([self _widthForSection:section] - ((_numberOfColumns - 1) * _columnSpacing)) / _numberOfColumns;
}

- (CGSize)itemSizeAtIndexPath:(NSIndexPath *)indexPath
{
    CGSize size = CGSizeMake([self _columnWidthForSection:indexPath.section], 0);
    CGSize originalSize = [[self _delegate] collectionView:self.collectionView layout:self originalItemSizeAtIndexPath:indexPath];
    if (originalSize.height > 0 && originalSize.width > 0) {
        size.height = originalSize.height / originalSize.width * size.width;
    }
    return size;
}

- (CGSize)headerSizeForSection:(NSInteger)section
{
    return CGSizeMake([self _widthForSection:section], _headerHeight);
}

- (CGSize)footerSizeForSection:(NSInteger)section {
    return CGSizeMake([self _widthForSection:section], _footerHeight);

}

- (CGSize)collectionViewContentSize
{
    CGFloat height = [[[_columnHeights lastObject] firstObject] floatValue] + _footerHeight;
    return CGSizeMake(self.collectionView.bounds.size.width, height);
}

- (NSUInteger)_tallestColumnIndexInSection:(NSUInteger)section
{
    __block NSUInteger index = 0;
    __block CGFloat tallestHeight = 0;
    [_columnHeights[section] enumerateObjectsUsingBlock:^(NSNumber *height, NSUInteger idx, BOOL *stop) {
        if (height.floatValue > tallestHeight) {
            index = idx;
            tallestHeight = height.floatValue;
        }
    }];
    return index;
}

- (NSUInteger)_shortestColumnIndexInSection:(NSUInteger)section
{
    __block NSUInteger index = 0;
    __block CGFloat shortestHeight = CGFLOAT_MAX;
    [_columnHeights[section] enumerateObjectsUsingBlock:^(NSNumber *height, NSUInteger idx, BOOL *stop) {
        if (height.floatValue < shortestHeight) {
            index = idx;
            shortestHeight = height.floatValue;
        }
    }];
    return index;
}

- (id<BACollectionViewLayoutDelegate>)_delegate
{
    return (id<BACollectionViewLayoutDelegate>)self.collectionView.delegate;
}

@end
